{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2812d42-da69-4741-97e8-5d1909931d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Optional\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "from datasets import load_dataset, load_metric\n",
    "\n",
    "import transformers\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    HfArgumentParser,\n",
    "    PretrainedConfig,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    default_data_collator,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import check_min_version, send_example_telemetry\n",
    "from transformers.utils.versions import require_version\n",
    "\n",
    "\n",
    "# Will error if the minimal version of Transformers is not installed. Remove at your own risks.\n",
    "check_min_version(\"4.21.0\")\n",
    "\n",
    "require_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/text-classification/requirements.txt\")\n",
    "\n",
    "task_to_keys = {\n",
    "    \"cola\": (\"sentence\", None),\n",
    "    \"mnli\": (\"premise\", \"hypothesis\"),\n",
    "    \"mrpc\": (\"sentence1\", \"sentence2\"),\n",
    "    \"qnli\": (\"question\", \"sentence\"),\n",
    "    \"qqp\": (\"question1\", \"question2\"),\n",
    "    \"rte\": (\"sentence1\", \"sentence2\"),\n",
    "    \"sst2\": (\"sentence\", None),\n",
    "    \"stsb\": (\"sentence1\", \"sentence2\"),\n",
    "    \"wnli\": (\"sentence1\", \"sentence2\"),\n",
    "}\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "\n",
    "    Using `HfArgumentParser` we can turn this class\n",
    "    into argparse arguments to be able to specify them on\n",
    "    the command line.\n",
    "    \"\"\"\n",
    "\n",
    "    task_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The name of the task to train on: \" + \", \".join(task_to_keys.keys())},\n",
    "    )\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    dataset_config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The configuration name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    max_seq_length: int = field(\n",
    "        default=128,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"The maximum total input sequence length after tokenization. Sequences longer \"\n",
    "                \"than this will be truncated, sequences shorter will be padded.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached preprocessed datasets or not.\"}\n",
    "    )\n",
    "    pad_to_max_length: bool = field(\n",
    "        default=True,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"Whether to pad all samples to `max_seq_length`. \"\n",
    "                \"If False, will pad the samples dynamically when batching to the maximum length in the batch.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    max_train_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of training examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    max_eval_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of evaluation examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    max_predict_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of prediction examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    train_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the training data.\"}\n",
    "    )\n",
    "    validation_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the validation data.\"}\n",
    "    )\n",
    "    test_file: Optional[str] = field(default=None, metadata={\"help\": \"A csv or a json file containing the test data.\"})\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.task_name is not None:\n",
    "            self.task_name = self.task_name.lower()\n",
    "            if self.task_name not in task_to_keys.keys():\n",
    "                raise ValueError(\"Unknown task, you should pick one in \" + \",\".join(task_to_keys.keys()))\n",
    "        elif self.dataset_name is not None:\n",
    "            pass\n",
    "        elif self.train_file is None or self.validation_file is None:\n",
    "            raise ValueError(\"Need either a GLUE task, a training/validation file or a dataset name.\")\n",
    "        else:\n",
    "            train_extension = self.train_file.split(\".\")[-1]\n",
    "            assert train_extension in [\"csv\", \"json\"], \"`train_file` should be a csv or a json file.\"\n",
    "            validation_extension = self.validation_file.split(\".\")[-1]\n",
    "            assert (\n",
    "                validation_extension == train_extension\n",
    "            ), \"`validation_file` should have the same extension (csv or json) as `train_file`.\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ModelArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.\n",
    "    \"\"\"\n",
    "\n",
    "    model_name_or_path: str = field(\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
    "    )\n",
    "    config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"Pretrained config name or path if not the same as model_name\"}\n",
    "    )\n",
    "    tokenizer_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"Pretrained tokenizer name or path if not the same as model_name\"}\n",
    "    )\n",
    "    cache_dir: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Where do you want to store the pretrained models downloaded from huggingface.co\"},\n",
    "    )\n",
    "    use_fast_tokenizer: bool = field(\n",
    "        default=True,\n",
    "        metadata={\"help\": \"Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.\"},\n",
    "    )\n",
    "    model_revision: str = field(\n",
    "        default=\"main\",\n",
    "        metadata={\"help\": \"The specific model version to use (can be a branch name, tag name or commit id).\"},\n",
    "    )\n",
    "    use_auth_token: bool = field(\n",
    "        default=False,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"Will use the token generated when running `transformers-cli login` (necessary to use this script \"\n",
    "                \"with private models).\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    ignore_mismatched_sizes: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Will enable to load a pretrained model whose head dimensions are different.\"},\n",
    "    )\n",
    "\n",
    "\n",
    "def main():\n",
    "    # See all possible arguments in src/transformers/training_args.py\n",
    "    # or by passing the --help flag to this script.\n",
    "    # We now keep distinct sets of args, for a cleaner separation of concerns.\n",
    "\n",
    "    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))\n",
    "    if len(sys.argv) == 2 and sys.argv[1].endswith(\".json\"):\n",
    "        # If we pass only one argument to the script and it's the path to a json file,\n",
    "        # let's parse it to get our arguments.\n",
    "        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))\n",
    "    else:\n",
    "        model_args, data_args, training_args = parser.parse_args_into_dataclasses()\n",
    "\n",
    "    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The\n",
    "    # information sent is the one passed as arguments along with your Python/PyTorch versions.\n",
    "    send_example_telemetry(\"run_glue\", model_args, data_args)\n",
    "\n",
    "    # Setup logging\n",
    "    logging.basicConfig(\n",
    "        format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
    "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "        handlers=[logging.StreamHandler(sys.stdout)],\n",
    "    )\n",
    "\n",
    "    log_level = training_args.get_process_log_level()\n",
    "    logger.setLevel(log_level)\n",
    "    datasets.utils.logging.set_verbosity(log_level)\n",
    "    transformers.utils.logging.set_verbosity(log_level)\n",
    "    transformers.utils.logging.enable_default_handler()\n",
    "    transformers.utils.logging.enable_explicit_format()\n",
    "\n",
    "    # Log on each process the small summary:\n",
    "    logger.warning(\n",
    "        f\"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\"\n",
    "        + f\"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}\"\n",
    "    )\n",
    "    logger.info(f\"Training/evaluation parameters {training_args}\")\n",
    "\n",
    "    # Detecting last checkpoint.\n",
    "    last_checkpoint = None\n",
    "    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:\n",
    "        last_checkpoint = get_last_checkpoint(training_args.output_dir)\n",
    "        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:\n",
    "            raise ValueError(\n",
    "                f\"Output directory ({training_args.output_dir}) already exists and is not empty. \"\n",
    "                \"Use --overwrite_output_dir to overcome.\"\n",
    "            )\n",
    "        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:\n",
    "            logger.info(\n",
    "                f\"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change \"\n",
    "                \"the `--output_dir` or add `--overwrite_output_dir` to train from scratch.\"\n",
    "            )\n",
    "\n",
    "    # Set seed before initializing model.\n",
    "    set_seed(training_args.seed)\n",
    "\n",
    "    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)\n",
    "    # or specify a GLUE benchmark task (the dataset will be downloaded automatically from the datasets Hub).\n",
    "    #\n",
    "    # For CSV/JSON files, this script will use as labels the column called 'label' and as pair of sentences the\n",
    "    # sentences in columns called 'sentence1' and 'sentence2' if such column exists or the first two columns not named\n",
    "    # label if at least two columns are provided.\n",
    "    #\n",
    "    # If the CSVs/JSONs contain only one non-label column, the script does single sentence classification on this\n",
    "    # single column. You can easily tweak this behavior (see below)\n",
    "    #\n",
    "    # In distributed training, the load_dataset function guarantee that only one local process can concurrently\n",
    "    # download the dataset.\n",
    "    if data_args.task_name is not None:\n",
    "        # Downloading and loading a dataset from the hub.\n",
    "        raw_datasets = load_dataset(\n",
    "            \"glue\",\n",
    "            data_args.task_name,\n",
    "            cache_dir=model_args.cache_dir,\n",
    "            use_auth_token=True if model_args.use_auth_token else None,\n",
    "        )\n",
    "    elif data_args.dataset_name is not None:\n",
    "        # Downloading and loading a dataset from the hub.\n",
    "        raw_datasets = load_dataset(\n",
    "            data_args.dataset_name,\n",
    "            data_args.dataset_config_name,\n",
    "            cache_dir=model_args.cache_dir,\n",
    "            use_auth_token=True if model_args.use_auth_token else None,\n",
    "        )\n",
    "    else:\n",
    "        # Loading a dataset from your local files.\n",
    "        # CSV/JSON training and evaluation files are needed.\n",
    "        data_files = {\"train\": data_args.train_file, \"validation\": data_args.validation_file}\n",
    "\n",
    "        # Get the test dataset: you can provide your own CSV/JSON test file (see below)\n",
    "        # when you use `do_predict` without specifying a GLUE benchmark task.\n",
    "        if training_args.do_predict:\n",
    "            if data_args.test_file is not None:\n",
    "                train_extension = data_args.train_file.split(\".\")[-1]\n",
    "                test_extension = data_args.test_file.split(\".\")[-1]\n",
    "                assert (\n",
    "                    test_extension == train_extension\n",
    "                ), \"`test_file` should have the same extension (csv or json) as `train_file`.\"\n",
    "                data_files[\"test\"] = data_args.test_file\n",
    "            else:\n",
    "                raise ValueError(\"Need either a GLUE task or a test file for `do_predict`.\")\n",
    "\n",
    "        for key in data_files.keys():\n",
    "            logger.info(f\"load a local file for {key}: {data_files[key]}\")\n",
    "\n",
    "        if data_args.train_file.endswith(\".csv\"):\n",
    "            # Loading a dataset from local csv files\n",
    "            raw_datasets = load_dataset(\n",
    "                \"csv\",\n",
    "                data_files=data_files,\n",
    "                cache_dir=model_args.cache_dir,\n",
    "                use_auth_token=True if model_args.use_auth_token else None,\n",
    "            )\n",
    "        else:\n",
    "            # Loading a dataset from local json files\n",
    "            raw_datasets = load_dataset(\n",
    "                \"json\",\n",
    "                data_files=data_files,\n",
    "                cache_dir=model_args.cache_dir,\n",
    "                use_auth_token=True if model_args.use_auth_token else None,\n",
    "            )\n",
    "    # See more about loading any type of standard or custom dataset at\n",
    "    # https://huggingface.co/docs/datasets/loading_datasets.html.\n",
    "\n",
    "    # Labels\n",
    "    if data_args.task_name is not None:\n",
    "        is_regression = data_args.task_name == \"stsb\"\n",
    "        if not is_regression:\n",
    "            label_list = raw_datasets[\"train\"].features[\"label\"].names\n",
    "            num_labels = len(label_list)\n",
    "        else:\n",
    "            num_labels = 1\n",
    "    else:\n",
    "        # Trying to have good defaults here, don't hesitate to tweak to your needs.\n",
    "        is_regression = raw_datasets[\"train\"].features[\"label\"].dtype in [\"float32\", \"float64\"]\n",
    "        if is_regression:\n",
    "            num_labels = 1\n",
    "        else:\n",
    "            # A useful fast method:\n",
    "            # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.unique\n",
    "            label_list = raw_datasets[\"train\"].unique(\"label\")\n",
    "            label_list.sort()  # Let's sort it for determinism\n",
    "            num_labels = len(label_list)\n",
    "\n",
    "    # Load pretrained model and tokenizer\n",
    "    #\n",
    "    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n",
    "    # download model & vocab.\n",
    "    config = AutoConfig.from_pretrained(\n",
    "        model_args.config_name if model_args.config_name else model_args.model_name_or_path,\n",
    "        num_labels=num_labels,\n",
    "        finetuning_task=data_args.task_name,\n",
    "        cache_dir=model_args.cache_dir,\n",
    "        revision=model_args.model_revision,\n",
    "        use_auth_token=True if model_args.use_auth_token else None,\n",
    "    )\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\n",
    "        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,\n",
    "        cache_dir=model_args.cache_dir,\n",
    "        use_fast=model_args.use_fast_tokenizer,\n",
    "        revision=model_args.model_revision,\n",
    "        use_auth_token=True if model_args.use_auth_token else None,\n",
    "    )\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        model_args.model_name_or_path,\n",
    "        from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "        config=config,\n",
    "        cache_dir=model_args.cache_dir,\n",
    "        revision=model_args.model_revision,\n",
    "        use_auth_token=True if model_args.use_auth_token else None,\n",
    "        ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,\n",
    "    )\n",
    "\n",
    "    # Preprocessing the raw_datasets\n",
    "    if data_args.task_name is not None:\n",
    "        sentence1_key, sentence2_key = task_to_keys[data_args.task_name]\n",
    "    else:\n",
    "        # Again, we try to have some nice defaults but don't hesitate to tweak to your use case.\n",
    "        non_label_column_names = [name for name in raw_datasets[\"train\"].column_names if name != \"label\"]\n",
    "        if \"sentence1\" in non_label_column_names and \"sentence2\" in non_label_column_names:\n",
    "            sentence1_key, sentence2_key = \"sentence1\", \"sentence2\"\n",
    "        else:\n",
    "            if len(non_label_column_names) >= 2:\n",
    "                sentence1_key, sentence2_key = non_label_column_names[:2]\n",
    "            else:\n",
    "                sentence1_key, sentence2_key = non_label_column_names[0], None\n",
    "\n",
    "    # Padding strategy\n",
    "    if data_args.pad_to_max_length:\n",
    "        padding = \"max_length\"\n",
    "    else:\n",
    "        # We will pad later, dynamically at batch creation, to the max sequence length in each batch\n",
    "        padding = False\n",
    "\n",
    "    # Some models have set the order of the labels to use, so let's make sure we do use it.\n",
    "    label_to_id = None\n",
    "    if (\n",
    "        model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id\n",
    "        and data_args.task_name is not None\n",
    "        and not is_regression\n",
    "    ):\n",
    "        # Some have all caps in their config, some don't.\n",
    "        label_name_to_id = {k.lower(): v for k, v in model.config.label2id.items()}\n",
    "        if list(sorted(label_name_to_id.keys())) == list(sorted(label_list)):\n",
    "            label_to_id = {i: int(label_name_to_id[label_list[i]]) for i in range(num_labels)}\n",
    "        else:\n",
    "            logger.warning(\n",
    "                \"Your model seems to have been trained with labels, but they don't match the dataset: \",\n",
    "                f\"model labels: {list(sorted(label_name_to_id.keys()))}, dataset labels: {list(sorted(label_list))}.\"\n",
    "                \"\\nIgnoring the model labels as a result.\",\n",
    "            )\n",
    "    elif data_args.task_name is None and not is_regression:\n",
    "        label_to_id = {v: i for i, v in enumerate(label_list)}\n",
    "\n",
    "    if label_to_id is not None:\n",
    "        model.config.label2id = label_to_id\n",
    "        model.config.id2label = {id: label for label, id in config.label2id.items()}\n",
    "    elif data_args.task_name is not None and not is_regression:\n",
    "        model.config.label2id = {l: i for i, l in enumerate(label_list)}\n",
    "        model.config.id2label = {id: label for label, id in config.label2id.items()}\n",
    "\n",
    "    if data_args.max_seq_length > tokenizer.model_max_length:\n",
    "        logger.warning(\n",
    "            f\"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the\"\n",
    "            f\"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}.\"\n",
    "        )\n",
    "    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)\n",
    "\n",
    "    def preprocess_function(examples):\n",
    "        # Tokenize the texts\n",
    "        args = (\n",
    "            (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])\n",
    "        )\n",
    "        result = tokenizer(*args, padding=padding, max_length=max_seq_length, truncation=True)\n",
    "\n",
    "        # Map labels to IDs (not necessary for GLUE tasks)\n",
    "        if label_to_id is not None and \"label\" in examples:\n",
    "            result[\"label\"] = [(label_to_id[l] if l != -1 else -1) for l in examples[\"label\"]]\n",
    "        return result\n",
    "\n",
    "    with training_args.main_process_first(desc=\"dataset map pre-processing\"):\n",
    "        raw_datasets = raw_datasets.map(\n",
    "            preprocess_function,\n",
    "            batched=True,\n",
    "            load_from_cache_file=not data_args.overwrite_cache,\n",
    "            desc=\"Running tokenizer on dataset\",\n",
    "        )\n",
    "    if training_args.do_train:\n",
    "        if \"train\" not in raw_datasets:\n",
    "            raise ValueError(\"--do_train requires a train dataset\")\n",
    "        train_dataset = raw_datasets[\"train\"]\n",
    "        if data_args.max_train_samples is not None:\n",
    "            max_train_samples = min(len(train_dataset), data_args.max_train_samples)\n",
    "            train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "    if training_args.do_eval:\n",
    "        if \"validation\" not in raw_datasets and \"validation_matched\" not in raw_datasets:\n",
    "            raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "        eval_dataset = raw_datasets[\"validation_matched\" if data_args.task_name == \"mnli\" else \"validation\"]\n",
    "        if data_args.max_eval_samples is not None:\n",
    "            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)\n",
    "            eval_dataset = eval_dataset.select(range(max_eval_samples))\n",
    "\n",
    "    if training_args.do_predict or data_args.task_name is not None or data_args.test_file is not None:\n",
    "        if \"test\" not in raw_datasets and \"test_matched\" not in raw_datasets:\n",
    "            raise ValueError(\"--do_predict requires a test dataset\")\n",
    "        predict_dataset = raw_datasets[\"test_matched\" if data_args.task_name == \"mnli\" else \"test\"]\n",
    "        if data_args.max_predict_samples is not None:\n",
    "            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)\n",
    "            predict_dataset = predict_dataset.select(range(max_predict_samples))\n",
    "\n",
    "    # Log a few random samples from the training set:\n",
    "    if training_args.do_train:\n",
    "        for index in random.sample(range(len(train_dataset)), 3):\n",
    "            logger.info(f\"Sample {index} of the training set: {train_dataset[index]}.\")\n",
    "\n",
    "    # Get the metric function\n",
    "    if data_args.task_name is not None:\n",
    "        metric = load_metric(\"glue\", data_args.task_name)\n",
    "    else:\n",
    "        metric = load_metric(\"accuracy\")\n",
    "\n",
    "    # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a\n",
    "    # predictions and label_ids field) and has to return a dictionary string to float.\n",
    "    def compute_metrics(p: EvalPrediction):\n",
    "        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "        preds = np.squeeze(preds) if is_regression else np.argmax(preds, axis=1)\n",
    "        if data_args.task_name is not None:\n",
    "            result = metric.compute(predictions=preds, references=p.label_ids)\n",
    "            if len(result) > 1:\n",
    "                result[\"combined_score\"] = np.mean(list(result.values())).item()\n",
    "            return result\n",
    "        elif is_regression:\n",
    "            return {\"mse\": ((preds - p.label_ids) ** 2).mean().item()}\n",
    "        else:\n",
    "            return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}\n",
    "\n",
    "    # Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if\n",
    "    # we already did the padding.\n",
    "    if data_args.pad_to_max_length:\n",
    "        data_collator = default_data_collator\n",
    "    elif training_args.fp16:\n",
    "        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)\n",
    "    else:\n",
    "        data_collator = None\n",
    "\n",
    "    # Initialize our Trainer\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=train_dataset if training_args.do_train else None,\n",
    "        eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "        compute_metrics=compute_metrics,\n",
    "        tokenizer=tokenizer,\n",
    "        data_collator=data_collator,\n",
    "    )\n",
    "   \n",
    "    # Training\n",
    "    if training_args.do_train:\n",
    "        checkpoint = None\n",
    "        if training_args.resume_from_checkpoint is not None:\n",
    "            checkpoint = training_args.resume_from_checkpoint\n",
    "        elif last_checkpoint is not None:\n",
    "            checkpoint = last_checkpoint\n",
    "        train_result = trainer.train(resume_from_checkpoint=checkpoint)\n",
    "        metrics = train_result.metrics\n",
    "        max_train_samples = (\n",
    "            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)\n",
    "        )\n",
    "        metrics[\"train_samples\"] = min(max_train_samples, len(train_dataset))\n",
    "\n",
    "        trainer.save_model()  # Saves the tokenizer too for easy upload\n",
    "\n",
    "        trainer.log_metrics(\"train\", metrics)\n",
    "        trainer.save_metrics(\"train\", metrics)\n",
    "        trainer.save_state()\n",
    "\n",
    "    # Evaluation\n",
    "    if training_args.do_eval:\n",
    "        logger.info(\"*** Evaluate ***\")\n",
    "\n",
    "        # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "        tasks = [data_args.task_name]\n",
    "        eval_datasets = [eval_dataset]\n",
    "        if data_args.task_name == \"mnli\":\n",
    "            tasks.append(\"mnli-mm\")\n",
    "            eval_datasets.append(raw_datasets[\"validation_mismatched\"])\n",
    "            combined = {}\n",
    "\n",
    "        for eval_dataset, task in zip(eval_datasets, tasks):\n",
    "            metrics = trainer.evaluate(eval_dataset=eval_dataset)\n",
    "\n",
    "            max_eval_samples = (\n",
    "                data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)\n",
    "            )\n",
    "            metrics[\"eval_samples\"] = min(max_eval_samples, len(eval_dataset))\n",
    "\n",
    "            if task == \"mnli-mm\":\n",
    "                metrics = {k + \"_mm\": v for k, v in metrics.items()}\n",
    "            if task is not None and \"mnli\" in task:\n",
    "                combined.update(metrics)\n",
    "\n",
    "            trainer.log_metrics(\"eval\", metrics)\n",
    "            trainer.save_metrics(\"eval\", combined if task is not None and \"mnli\" in task else metrics)\n",
    "\n",
    "    if training_args.do_predict:\n",
    "        logger.info(\"*** Predict ***\")\n",
    "\n",
    "        # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "        tasks = [data_args.task_name]\n",
    "        predict_datasets = [predict_dataset]\n",
    "        if data_args.task_name == \"mnli\":\n",
    "            tasks.append(\"mnli-mm\")\n",
    "            predict_datasets.append(raw_datasets[\"test_mismatched\"])\n",
    "\n",
    "        for predict_dataset, task in zip(predict_datasets, tasks):\n",
    "            # Removing the `label` columns because it contains -1 and Trainer won't like that.\n",
    "            predict_dataset = predict_dataset.remove_columns(\"label\")\n",
    "            predictions = trainer.predict(predict_dataset, metric_key_prefix=\"predict\").predictions\n",
    "            predictions = np.squeeze(predictions) if is_regression else np.argmax(predictions, axis=1)\n",
    "\n",
    "            output_predict_file = os.path.join(training_args.output_dir, f\"predict_results_{task}.txt\")\n",
    "            if trainer.is_world_process_zero():\n",
    "                with open(output_predict_file, \"w\") as writer:\n",
    "                    logger.info(f\"***** Predict results {task} *****\")\n",
    "                    writer.write(\"index\\tprediction\\n\")\n",
    "                    for index, item in enumerate(predictions):\n",
    "                        if is_regression:\n",
    "                            writer.write(f\"{index}\\t{item:3.3f}\\n\")\n",
    "                        else:\n",
    "                            item = label_list[item]\n",
    "                            writer.write(f\"{index}\\t{item}\\n\")\n",
    "\n",
    "    kwargs = {\"finetuned_from\": model_args.model_name_or_path, \"tasks\": \"text-classification\"}\n",
    "    if data_args.task_name is not None:\n",
    "        kwargs[\"language\"] = \"en\"\n",
    "        kwargs[\"dataset_tags\"] = \"glue\"\n",
    "        kwargs[\"dataset_args\"] = data_args.task_name\n",
    "        kwargs[\"dataset\"] = f\"GLUE {data_args.task_name.upper()}\"\n",
    "\n",
    "    if training_args.push_to_hub:\n",
    "        trainer.push_to_hub(**kwargs)\n",
    "    else:\n",
    "        trainer.create_model_card(**kwargs)\n",
    "\n",
    "\n",
    "def _mp_fn(index):\n",
    "    # For xla_spawn (TPUs)\n",
    "    main()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
